import decimal
import logging

import inflect

from app.migrate_neo4j.base import BaseMigrator

logger = logging.getLogger(__name__)


# -----------------------------------------------------------------------------
# --- 默认的、基于规则的 Migrator ---
# -----------------------------------------------------------------------------
class DefaultRelsToGraphMigrator(BaseMigrator):
    """
    提供基于严格外键约束的建模和转换逻辑。
    适用于数据库设计严谨的场景。
    """

    def __init__(self, rdbms_config, neo4j_config, batch_size=1000):
        super().__init__(rdbms_config, neo4j_config, batch_size)
        self.p = inflect.engine()

    def model_schema(self):
        logging.info("[MODEL] Modeling schema based on foreign key rules.")
        model = {"nodes": {}, "relationships": []}
        table_names = self.inspector.get_table_names()
        for table_name in table_names:
            pk = self.inspector.get_pk_constraint(table_name)['constrained_columns'][0]
            label = self.p.singular_noun(table_name) or table_name
            model["nodes"][table_name] = {"label_name": label.capitalize(), "primary_key": pk}

        for table_name in table_names:
            for fk in self.inspector.get_foreign_keys(table_name):
                start_table = fk['referred_table']
                end_table = table_name
                model["relationships"].append({
                    "source_table_name": end_table,
                    "start_node_label": model["nodes"][start_table]["label_name"],
                    "start_node_key": fk['referred_columns'][0],
                    "end_node_label": model["nodes"][end_table]["label_name"],
                    "end_node_key": model["nodes"][end_table]["primary_key"],
                    "rel_type": f"HAS_{model['nodes'][end_table]['label_name'].upper()}",
                    "fk_column": fk['constrained_columns'][0]
                })
        self.graph_model = model

    def transform_data(self, data_row, context):
        if context['type'] == 'node':
            for key, value in data_row.items():
                if isinstance(value, decimal.Decimal):
                    data_row[key] = float(value)
            return data_row
        elif context['type'] == 'relationship':
            fk_val = data_row.get(context['fk_column'])
            pk_val = data_row.get(context['end_node_key'])
            if fk_val is not None and pk_val is not None:
                return {"start_id": fk_val, "end_id": pk_val}
        return None
